import numpy
import frame_utils
        
def sta(frames, frame_times, spike_times, tau, num_samples, shade=1.0):
    """
    Calculate the spike triggered average.
    Inputs:
        frames      : A 3D numpy array of formated frames. (see format_frames)
        frame_times : A list of times when the frames were drawn.
        spike_times : A list of times when 'events' or spikes occured
        tau         : the window of time BEFORE an 'event' to look at the stimulus.
        num_samples : sets the temporal resolution of the sta.
        shade=1.0       : The window background shade prior to stimulation.
    Returns:
        sta         : A 3D numpy array of formated frames, this time 
                       .shape[0]==num_samples
    """
    if type(frames) == list:
        frames = frame_utils.format_frames(frames)
       
    tau_frames = []
    for st in spike_times:
        tchunk =  frame_utils.stim_chunk(st, tau, num_samples, frames, frame_times) 
        if tchunk is not None:
            tau_frames.append(tchunk)
    
    tau_frames = numpy.array(tau_frames)
    return_value = numpy.average(tau_frames, axis=0)
    return return_value, len(tau_frames)

